import os
import numpy as np
from scipy.spatial.transform  import Rotation as sRot
import torch


def rot_decompose(R, axis):
    """
    R: rotation matrices, shape = (N, 3, 3)
    axis: shape = (3)

    return: (N, )
    """
    axis2 = sRot.from_matrix(R).apply(axis)
    axis = torch.from_numpy(axis)
    axis2 = torch.from_numpy(axis2)
    theta = torch.arccos(torch.sum(axis.unsqueeze(0) * axis2, dim=1).clamp(-1, 1))
    axis3 = torch.cross(axis2, axis.unsqueeze(0), dim=1)
    Rp = (axis3 / torch.norm(axis3, dim=1).unsqueeze(1).clamp(1e-6, None)) * theta.unsqueeze(1)
    Ry = sRot.from_rotvec(Rp) * sRot.from_matrix(R)
    Rxz = (Ry.inv() * sRot.from_matrix(R)).as_matrix()
    Ry_theta = torch.sum(torch.from_numpy(Ry.as_rotvec()) * axis, dim=1)
    return Ry_theta, Rxz
